#include <iostream>
#include <hip/hip_runtime.h>

#define HIP_CHECK(cmd) \
do { \
    hipError_t error = cmd; \
    if (error != hipSuccess) { \
        std::cerr << "HIP error: " << hipGetErrorString(error) \
                  << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    } \
} while(0)

__global__ void vectorAdd(const float* A, const float* B, float* C, int numElements) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < numElements) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    const int N = 1 << 20;  // 1M elements
    const size_t bytes = N * sizeof(float);

    // 分配主机内存
    float *h_A = new float[N];
    float *h_B = new float[N];
    float *h_C = new float[N];

    // 初始化主机数据
    for (int i = 0; i < N; ++i) {
        h_A[i] = 1.0f;
        h_B[i] = 2.0f;
    }

    // 分配设备内存
    float *d_A, *d_B, *d_C;
    HIP_CHECK(hipMalloc(&d_A, bytes));
    HIP_CHECK(hipMalloc(&d_B, bytes));
    HIP_CHECK(hipMalloc(&d_C, bytes));

    // 创建流
    hipStream_t stream1, stream2, stream3;
    HIP_CHECK(hipStreamCreate(&stream1));

    unsigned int flags = hipStreamNonBlocking;
    HIP_CHECK(hipStreamCreateWithFlags(&stream2, flags));

    int priority;
    HIP_CHECK(hipDeviceGetStreamPriorityRange(nullptr, &priority));
    HIP_CHECK(hipStreamCreateWithPriority(&stream3, flags, priority));

    // 创建事件
    hipEvent_t start, stop;
    HIP_CHECK(hipEventCreate(&start));
    HIP_CHECK(hipEventCreateWithFlags(&stop, hipEventDefault));

    // 记录开始事件
    HIP_CHECK(hipEventRecord(start, stream1));

    // 异步数据传输
    HIP_CHECK(hipMemcpyAsync(d_A, h_A, bytes, hipMemcpyHostToDevice, stream1));
    HIP_CHECK(hipMemcpyAsync(d_B, h_B, bytes, hipMemcpyHostToDevice, stream1));

    // 定义内核配置
    const int blockSize = 256;
    const int gridSize = (N + blockSize - 1) / blockSize;

    // 异步执行内核
//    vectorAdd<<<gridSize, blockSize, 0, stream1>>>(d_A, d_B, d_C);
    hipLaunchKernelGGL(vectorAdd, dim3(gridSize), dim3(blockSize), 0, 0, d_A, d_B, d_C, N);
    HIP_CHECK(hipGetLastError());

    // 异步数据传输回主机
    HIP_CHECK(hipMemcpyAsync(h_C, d_C, bytes, hipMemcpyDeviceToHost, stream1));

    // 记录结束事件
    HIP_CHECK(hipEventRecord(stop, stream1));

    // 等待事件完成
    HIP_CHECK(hipEventSynchronize(stop));

    // 计算事件之间的时间差
    float milliseconds = 0;
    HIP_CHECK(hipEventElapsedTime(&milliseconds, start, stop));
    std::cout << "Kernel execution time: " << milliseconds << " ms" << std::endl;

    // 验证结果
    for (int i = 0; i < N; ++i) {
        if (h_C[i] != 3.0f) {
            std::cerr << "Verification failed at index " << i << std::endl;
            return EXIT_FAILURE;
        }
    }

//    // 销毁事件
//    HIP_CHECK(hipEventDestroy(start));
//    HIP_CHECK(hipEventDestroy(stop));
//
//    // 销毁流
//    HIP_CHECK(hipStreamDestroy(stream1));
//    HIP_CHECK(hipStreamDestroy(stream2));
//    HIP_CHECK(hipStreamDestroy(stream3));

    // 释放设备内存
    HIP_CHECK(hipFree(d_A));
    HIP_CHECK(hipFree(d_B));
    HIP_CHECK(hipFree(d_C));

    // 释放主机内存
    delete[] h_A;
    delete[] h_B;
    delete[] h_C;

    return EXIT_SUCCESS;
}